# 新增文件: ensumble_webapp.py
import os
import joblib
import numpy as np
import cv2
import faiss
import logging

# Gradio 可能在新版 API 上有细微差别，Blocks 或 Interface 均可用
import gradio as gr

logging.basicConfig(level=logging.INFO)
MODEL_PATH = os.path.join('.', 'saved_models', 'best_ensemble.pkl')

# 加载模型
if not os.path.exists(MODEL_PATH):
    logging.warning(f"未找到模型：{MODEL_PATH}。请先运行 train_ensemble.py 并将模型保存到 saved_models/.")
    model = None
else:
    model = joblib.load(MODEL_PATH)
    logging.info(f"已加载模型: {MODEL_PATH}")


def preprocess_image_pil(pil_img):
    # 接收 PIL.Image，转为灰度，调整为 32x32，展平并转为 float32
    img = np.array(pil_img.convert('L'))
    img = cv2.resize(img, (32, 32))
    x = img.reshape(1, -1).astype('float32')
    try:
        faiss.normalize_L2(x)
    except Exception:
        pass
    return x


def predict(pil_img):
    if model is None:
        return '模型未加载，请先训练并保存模型。'
    x = preprocess_image_pil(pil_img)
    # 兼容 sklearn 管线和自定义 FaissKNeighbors
    try:
        if hasattr(model, 'predict_proba'):
            proba = model.predict_proba(x)
            idx = int(np.argmax(proba, axis=1)[0])
            label = 'dog' if idx == 1 else 'cat'
            conf = float(proba[0, idx])
            return f"{label} (confidence: {conf:.3f})"
        else:
            pred = model.predict(x)
            idx = int(pred[0])
            label = 'dog' if idx == 1 else 'cat'
            return label
    except Exception as e:
        logging.exception('预测时出错')
        return f'预测失败: {e}'


def create_demo():
    title = '猫/狗 分类 - Ensemble 模型'
    description = '上传一张图片，模型将返回 cat 或 dog（使用 flat 特征）。确保已在 saved_models/ 下保存 best_ensemble.pkl。'
    iface = gr.Interface(fn=predict, inputs=gr.Image(type='pil'), outputs='text', title=title, description=description)
    return iface


if __name__ == '__main__':
    demo = create_demo()
    # launch 在本机运行，share=True 可以生成外网链接（需要注意 gradio 版本）
    demo.launch()